// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/output_stream.h>
#include <pollux/dwio/common/compression/compression.h>
#include <pollux/dwio/common/compression/compression_buffer_pool.h>

namespace kumo::pollux::dwio::common::compression {
    class PagedOutputStream : public BufferedOutputStream {
    public:
        PagedOutputStream(
            CompressionBufferPool &pool,
            DataBufferHolder &bufferHolder,
            uint32_t compressionThreshold,
            uint8_t pageHeaderSize,
            std::unique_ptr<Compressor> compressor,
            const dwio::common::encryption::Encrypter *encryptor)
            : BufferedOutputStream(bufferHolder),
              pool_{&pool},
              compressor_{std::move(compressor)},
              encryptor_{encryptor},
              threshold_{compressionThreshold},
              pageHeaderSize_{pageHeaderSize} {
            POLLUX_CHECK(
                compressor_ || encryptor_,
                "Neither compressor or encryptor is set for paged output stream");
        }

        bool Next(void **data, int32_t *size, uint64_t increment) override;

        uint64_t flush() override;

        uint64_t size() const override {
            // only care about compressed size
            return bufferHolder_.size();
        }

        void BackUp(int32_t count) override;

        std::string getName() const override {
            return "paged output stream";
        }

        void recordPosition(
            PositionRecorder &recorder,
            int32_t bufferLength,
            int32_t bufferOffset,
            int32_t strideIndex = -1) const override;

    private:
        // create page using compressor and encryptor
        std::vector<melon::StringPiece> createPage();

        void writeHeader(char *buffer, size_t compressedSize, bool original);

        void updateSize(char *buffer, size_t compressedSize);

        void resetBuffers();

        CompressionBufferPool *const pool_;

        const std::unique_ptr<Compressor> compressor_;

        // Encryption provider
        const dwio::common::encryption::Encrypter *const encryptor_;

        // threshold below which, we skip compression
        const uint32_t threshold_;

        const uint8_t pageHeaderSize_;

        // Buffer to hold compressed data
        std::unique_ptr<dwio::common::DataBuffer<char> > compressionBuffer_{nullptr};

        // buffer that holds encrypted data
        std::unique_ptr<melon::IOBuf> encryptionBuffer_{nullptr};
    };
} // namespace kumo::pollux::dwio::common::compression
